#!usr/bin/env python
# -*- coding:utf-8 _*-
""" 
@author:robot
@file: tap_connector05_postgres.py
@version:
@time: 2022/01/13 
"""
import psycopg2

from TapClientConnector.connector import connector_query_logger, generate_logger
from TapClientConnector.connector.tap_connector02_oracle import TAPOracleConnector
from db_info import Postgres


class TAPPostgresConnector(TAPOracleConnector):

    def __init__(self, kwargs):
        super(TAPOracleConnector, self).__init__(kwargs)
        self.connector = psycopg2.connect(**self.FORMAT_DB_INFO)

    @connector_query_logger
    def show_dbs(self) -> list:
        query_str = "select pg_database.datname, pg_database_size(pg_database.datname) AS size from pg_database"
        db_list = []
        for i in self._query(query_str):
            db_list.append(i[0])
        return db_list

    @connector_query_logger
    def show_tables(self) -> list:
        schema = self.DB_INFO.get('database_owner')
        query_str = f"select tablename from pg_tables where schemaname='{schema}'"
        table_list = []
        for i in self._query(query_str):
            table_list.append(i[0])
        return table_list

    @generate_logger
    def generate_alter_add(self, table, column, _type, other=''):
        if table:
            self.CURRENT_TABLE = table
        return f"""
                alter table "{table}" add column {column} {_type} {other}
                """

    @generate_logger
    def generate_alter_modify(self, table, column, _type):
        if table:
            self.CURRENT_TABLE = table
        return f"""
                alter table "{table}" alter column {column} type {_type}
                """

    @generate_logger
    def generate_alter_change(self, table, old_column, new_column, _type):
        if table:
            self.CURRENT_TABLE = table
        return f"""
                alter table "{table}" rename {old_column} to {new_column}
                """

    @generate_logger
    def generate_alter_drop(self, table, column):
        if table:
            self.CURRENT_TABLE = table
        return f"""
                alter table "{table}" drop column {column}
                """


if __name__ == '__main__':
    T = TAPPostgresConnector(Postgres.postgres_5437)
    print(T.query_get_column('select * from "pg2"'))
